"""
python p2_solution_pyomo.py
"""

from pyomo.environ import ConcreteModel, Var, Objective, Constraint, SolverFactory, Binary, NonNegativeReals
import json
import os

# 读取仓库数据
def read_data():
    with open('../fujian/fujian3/origin_data/warehouse.json') as f:
        warehouses = json.load(f)
    with open('../fujian/fujian3/data_from_p1/all_average_inventory.json') as f:
        category_inventory = {item['category_id']: item['average_inventory'] for item in json.load(f)}
    with open('../fujian/fujian3/data_from_p1/all_average_sales.json') as f:
        category_sales = {item['category_id']: item['average_sales'] for item in json.load(f)}
    return warehouses, category_inventory, category_sales

warehouses, category_inventory, category_sales = read_data()
warehouse_ids = [w['warehouse_id'] for w in warehouses]
category_ids = list(category_inventory.keys())

# 创建模型
model = ConcreteModel()

# 定义变量：x[c, w] 表示类别 c 是否分配到仓库 w（0 或 1）
model.x = Var(category_ids, warehouse_ids, domain=Binary)

# 定义目标函数：最小化总成本
def objective_rule(model):
    return sum(w['daily_cost'] * (sum(model.x[c, w['warehouse_id']] for c in category_ids) > 0) for w in warehouses)
model.objective = Objective(rule=objective_rule, sense=1)

# 约束1：每个类别只能分配到一个仓库
def category_constraint_rule(model, c):
    return sum(model.x[c, w] for w in warehouse_ids) == 1
model.category_constraint = Constraint(category_ids, rule=category_constraint_rule)

# 约束2：每个仓库的库存量小于其最大库存
def inventory_constraint_rule(model, w):
    warehouse = next(wh for wh in warehouses if wh['warehouse_id'] == w)
    return sum(model.x[c, w] * category_inventory[c] for c in category_ids) <= warehouse['max_inventory']
model.inventory_constraint = Constraint(warehouse_ids, rule=inventory_constraint_rule)

# 约束3：每个仓库的出货量小于其最大出货量
def sales_constraint_rule(model, w):
    warehouse = next(wh for wh in warehouses if wh['warehouse_id'] == w)
    return sum(model.x[c, w] * category_sales[c] for c in category_ids) <= warehouse['max_sales']
model.sales_constraint = Constraint(warehouse_ids, rule=sales_constraint_rule)

# 求解模型
solver = SolverFactory('glpk')
result = solver.solve(model)

# 处理结果
if result.solver.status == 'ok':
    # 分配结果
    assignment = [{"category_id": c, "warehouse_id": w} for c in category_ids for w in warehouse_ids if model.x[c, w].value == 1]

    # 输出分配结果到 JSON 文件
    output_dir = '../fujian/fujian3/Pyomo/'
    os.makedirs(output_dir, exist_ok=True)
    with open(os.path.join(output_dir, 'result.json'), 'w') as f:
        json.dump(assignment, f, indent=4)

    # 计算总成本
    total_cost = sum(w['daily_cost'] for w in warehouses if any(model.x[c, w['warehouse_id']].value == 1 for c in category_ids))
    with open(os.path.join(output_dir, 'all_cost.txt'), 'w') as f:
        f.write(str(total_cost))

    # 计算仓库的利用率
    inventory_utilization = []
    sales_utilization = []
    for w in warehouses:
        warehouse_id = w['warehouse_id']
        total_inventory = sum(model.x[c, warehouse_id].value * category_inventory[c] for c in category_ids)
        total_sales = sum(model.x[c, warehouse_id].value * category_sales[c] for c in category_ids)

        utilization_inventory = total_inventory / w['max_inventory'] if w['max_inventory'] > 0 else 0
        utilization_sales = total_sales / w['max_sales'] if w['max_sales'] > 0 else 0

        inventory_utilization.append({
            "warehouse_id": warehouse_id,
            "utilization_rate_of_inventory": utilization_inventory
        })
        sales_utilization.append({
            "warehouse_id": warehouse_id,
            "utilization_rate_of_sales": utilization_sales
        })

    # 输出库存利用率
    with open(os.path.join(output_dir, 'all_utilization_rate_inventory.json'), 'w') as f:
        json.dump(inventory_utilization, f, indent=4)

    # 输出出货量利用率
    with open(os.path.join(output_dir, 'all_utilization_rate_sales.json'), 'w') as f:
        json.dump(sales_utilization, f, indent=4)

else:
    print("未找到合适的解")
